# Function to create diagnostic plots
create_diagnostic_plots <- function(model, data) {
  cat("Creating diagnostic plots...\n")
  
  # Extract model diagnostics
  model_data <- data.frame(
    fitted = fitted(model),
    residuals = residuals(model),
    standardized_residuals = rstandard(model),
    sqrt_abs_std_residuals = sqrt(abs(rstandard(model))),
    leverage = hatvalues(model),
    cooks_distance = cooks.distance(model)
  )
  
  # Add observation numbers
  model_data$obs_number <- 1:nrow(model_data)
  
  # Define common theme (matching your existing theme)
  theme_diagnostic <- theme_minimal() +
    theme(
      plot.title = element_text(size = 12, face = "bold"),
      plot.subtitle = element_text(size = 10, color = "grey60"),
      axis.title = element_text(size = 10),
      axis.text = element_text(size = 9),
      legend.position = "none",
      panel.grid.minor = element_blank(),
      strip.text = element_text(size = 10, face = "bold")
    )
  
  # Plot 1: Residuals vs Fitted
  p1 <- ggplot(model_data, aes(x = fitted, y = residuals)) +
    geom_point(size = 2, alpha = 0.7, color = "#1F78B4") +
    geom_smooth(method = "loess", se = TRUE, color = "red", linetype = "dashed", alpha = 0.3) +
    geom_hline(yintercept = 0, linetype = "solid", color = "grey50", alpha = 0.7) +
    labs(
      title = "Residuals vs Fitted",
      subtitle = "Check for linearity and homoscedasticity",
      x = "Fitted Values",
      y = "Residuals"
    ) +
    theme_diagnostic
  
  # Plot 2: Normal Q-Q Plot
  p2 <- ggplot(model_data, aes(sample = standardized_residuals)) +
    stat_qq(size = 2, alpha = 0.7, color = "#1F78B4") +
    stat_qq_line(color = "red", linetype = "dashed") +
    labs(
      title = "Normal Q-Q",
      subtitle = "Check for normality of residuals",
      x = "Theoretical Quantiles",
      y = "Standardized Residuals"
    ) +
    theme_diagnostic
  
  # Plot 3: Scale-Location (Square root of standardized residuals vs fitted)
  p3 <- ggplot(model_data, aes(x = fitted, y = sqrt_abs_std_residuals)) +
    geom_point(size = 2, alpha = 0.7, color = "#1F78B4") +
    geom_smooth(method = "loess", se = TRUE, color = "red", linetype = "dashed", alpha = 0.3) +
    labs(
      title = "Scale-Location",
      subtitle = "Check for homoscedasticity",
      x = "Fitted Values",
      y = expression(sqrt("|Standardized Residuals|"))
    ) +
    theme_diagnostic
  
  # Plot 4: Residuals vs Leverage (Cook's distance)
  # Identify influential points (Cook's distance > 4/n)
  n <- nrow(model_data)
  cooks_threshold <- 4/n
  
  p4 <- model_data %>%
    # Add country names (assuming they're in final_data)
    mutate(country_name = final_data$country_name[1:nrow(model_data)]) %>%
    ggplot(aes(x = leverage, y = standardized_residuals)) +
    geom_point(aes(size = cooks_distance), alpha = 0.7, color = "#1F78B4") +
    geom_smooth(method = "loess", se = TRUE, color = "red", linetype = "dashed", alpha = 0.3) +
    geom_hline(yintercept = 0, linetype = "solid", color = "grey50", alpha = 0.7) +
    geom_hline(yintercept = c(-2, 2), linetype = "dashed", color = "grey70", alpha = 0.7) +
    # Label points with high Cook's distance
    ggrepel::geom_text_repel(
      aes(label = ifelse(cooks_distance > cooks_threshold, country_name, "")),
      size = 3, color = "red", alpha = 0.8
    ) +
    scale_size_continuous(range = c(1, 4)) +
    labs(
      title = "Residuals vs Leverage",
      subtitle = "Influential observations labeled (Cook's distance > 4/n)",
      x = "Leverage",
      y = "Standardized Residuals"
    ) +
    theme_diagnostic +
    theme(legend.position = "none")
  
  cat("...finished!\n")
  return(list(p1 = p1, p2 = p2, p3 = p3, p4 = p4))
}
